package main

import (
	"context"
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)

func NewLimiter(r rate.Limit, b int, t time.Duration) gin.HandlerFunc {
	limiters := &sync.Map{}
	return func(c *gin.Context) {
		// 获取限速器
		// key 除了 ip 之外， 也可以起是他，例如： header, user name 等
		key := c.ClientIP()
		l, _ := limiters.LoadOrStore(key, rate.NewLimiter(r, b))
		// 这里不要直接使用 gin 的 context， 默认是没有超时时间的
		ctx, cancel := context.WithTimeout(c, t)
		defer cancel()
		if err := l.(*rate.Limiter).Wait(ctx); err != nil {
			// 这里先不处理日志，如果返回错误就直接 429
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": err})
		}
		c.Next()
	}
}

func main() {
	e := gin.Default()
	// 新建一个限流器，允许突发 10 个并发，限速 3rps， 超过 500ms就不在等待
	e.Use(NewLimiter(3, 10, 500*time.Millisecond))
	e.GET("ping", func(c *gin.Context) {
		c.String(http.StatusOK, "pong")
	})
	e.Run(":8000")
}
